# -*- coding: utf-8 -*-
# !/usr/bin/python3
"""
Author :      wu
Description :
"""

import tensorflow as tf
from tensorflow.keras import layers, models, regularizers


# layers的子类需要重新实现初始化方法，build和call方法
class Linear(layers.Layer):
    def __init__(self, units=32, **kwargs):
        super(Linear, self).__init__(**kwargs)
        self.units = units

    # build方法一般设置需要训练的参数
    def build(self, input_shape):
        # 参数名是必须设置
        self.w = self.add_weight("w", shape=(input_shape[-1], self.units),
                                 initializer="random_normal",
                                 trainable=True)
        self.b = self.add_weight("b", shape=(self.units, ), initializer="random_normal",
                                 trainable=True)
        super(Linear, self).build(input_shape)  # 相当于设置self.built=True

    # call方法一般定义正向传播
    @tf.function
    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

    # 自定义的layer如果需要通过function模式组成模型时被保存为h5模型，需要定义get_config方法
    def get_config(self):
        config = super(Linear, self).get_config()
        config.update({"units": self.units})
        return config


def main():

    linear = Linear(units=1)
    linear.build(input_shape=(None, 2))

    tf.keras.backend.clear_session()
    model = models.Sequential()
    # model.add(linear)
    model.add(Linear(units=1, input_shape=(2, )))
    model.summary()
    print("model input shape: {}".format(model.input_shape))
    print("model output shape: {}".format(model.output_shape))

    model.save("./linear_layer.h5", save_format="h5")


if __name__ == "__main__":
    main()
